#include "../bulk_storage/bulk_storage.h"
#include "neighborhood_indexer.h"
#include "../query_handler/query_handler.h"
#include "../pdb_index/atom.h"
#include "aa_util.h"
#include "../m_tree/mTreeNode.h"

#include <string.h>
#include <sys/timeb.h>
#include <values.h>
#include <assert.h>

#include <stdio.h>
#include <errno.h>
#include <unistd.h>

void neighborhood_indexer_calculate_neighborhood(const char * pdbCode, 
						 double neighborhoodRadius,
						 off_t atomIndex,
						 struct store * atomStore, 
						 struct mTree * atomTree,
						 struct store * neighStore, 
						 struct mTree * neighTree,
						 void(*constructor)(const char* pdbCode, void * neigh,
								    struct store * atomStore, off_t center, off_t neighbors[], 
								    short count),
						 short neigh_count
						 );
void neighborhood_indexer_find_neighbors(struct mTree * tree, struct store * s, const off_t atomIndex, 
					 const double radius, const short neigh_count, off_t * neighbor_indices) ;


void openMainIndex(const char * neigh_base, struct store ** neighStore, struct mTree ** neighTree, size_t neigh_struct_size, 
		   double(*neigh_distance)(const void*,const void*), size_t store_cache_size, size_t tree_cache_size) {
  char * neigh_fat, *neigh_tree, *neigh_tree_fat;

  neigh_fat = malloc(sizeof(char) * (strlen(neigh_base) + 4 + 1));
  neigh_tree = malloc(sizeof(char) * (strlen(neigh_base) + 6 + 1));
  neigh_tree_fat = malloc(sizeof(char) * (strlen(neigh_base) + 10 + 1));

  strcpy(neigh_fat,      neigh_base);
  strcpy(neigh_tree,     neigh_base);
  strcpy(neigh_tree_fat, neigh_base);

  strcat(neigh_fat,      ".fat");
  strcat(neigh_tree,     ".mTree");
  strcat(neigh_tree_fat, ".mTree.fat");
  
  printf("Using these files for tree: %s, %s, %s, %s  sizeof(neighborhood) = %ld\n", neigh_base, neigh_fat, neigh_tree, neigh_tree_fat, (long)neigh_struct_size);
  fflush(stdout);
  
  *neighStore = store_open_cached(neigh_base, neigh_fat, neigh_struct_size, 100, store_cache_size);
  *neighTree  = mTree_open_cached(*neighStore, neigh_tree, neigh_tree_fat, tree_cache_size, neigh_distance);
  
  free(neigh_fat);
  free(neigh_tree);
  free(neigh_tree_fat);
}

int writeBackup(struct store * neighStore, struct mTree * neighTree, const char * atom_base, const char * backup_base, time_t * last_sync, time_t sync_interval) {
  time_t now;
  struct timeb ts;

  char * backup_fat;
  char * backup_tree;
  char * backup_tree_fat;
  
  backup_fat = malloc(sizeof(char) * (strlen(backup_base) + 4 + 1));
  backup_tree = malloc(sizeof(char) * (strlen(backup_base) + 6 + 1));
  backup_tree_fat = malloc(sizeof(char) * (strlen(backup_base) + 10 + 1));
  
  strcpy(backup_fat,      backup_base);
  strcpy(backup_tree,     backup_base);
  strcpy(backup_tree_fat, backup_base);
  
  strcat(backup_fat,      ".fat");
  strcat(backup_tree,     ".mTree");
  strcat(backup_tree_fat, ".mTree.fat");
  

  ftime(&ts);

  now = ts.time;

  if(now - *last_sync > sync_interval) {
    printf("Running disk sync and backup.\n");
    //      mTree_verify(neighTree);
    //      printf("Tree integrity verified.\n");
    printf("Skipping tree integrity check...\n");
    printf("Check for 'Backup complete.' after this line.\nIf it's missing, use the tilde copies of the backup files, and the\nrecovery_filename from the backup message before this one.\nrecovery_filename = %s\n",
	   atom_base);
    fflush(stdout);

    fsync(((struct cacheStore *)neighTree->nStore->store_p)->fileStore->fat_file);
    fsync(((struct cacheStore *)neighTree->nStore->store_p)->fileStore->data_file);
    fsync(((struct cacheStore *)neighTree->oStore->store_p)->fileStore->fat_file);
    fsync(((struct cacheStore *)neighTree->oStore->store_p)->fileStore->data_file);

    printf("Flushed file buffers of working copy to disk.\n");

    if(! (cs_prepareForBackup(neighTree->nStore->store_p, backup_tree, backup_tree_fat) && 
	  cs_prepareForBackup(neighTree->oStore->store_p, backup_base, backup_fat))) { 

      printf("Backup failed -- Couldn't move old backups out of the way.\n");
      printf("Check the time stamps to figure out what set of ~ files you've been left with.\n");
      
      fflush(stdout);

      return 0;

      
    }    


    //    if(! (cs_dumpBackup(neighTree->nStore->store_p, "node_backup", "node_backup.fat") && 
    //	  cs_dumpBackup(neighTree->oStore->store_p, "object_backup", "object_backup.fat"))) {
    if(! (cs_dumpBackup(neighTree->nStore->store_p, backup_tree, backup_tree_fat) && 
	  cs_dumpBackup(neighTree->oStore->store_p, backup_base, backup_fat))) { 
      printf("Backup failed!\n");
      printf("Attempting clean shutdown.");
      fflush(stdout);
      mTree_close(neighTree);
      store_close(neighStore);
      printf("Shutdown complete-If the line above this one reads 'Attempting clean shutdown',  then everything should be OK.\n");
      fflush(stdout);
      return 0;
    }

    printf("Backup complete.\n");
    fflush(stdout);
    ftime(&ts);
    *last_sync = ts.time;

  }

  free( backup_fat );
  free( backup_tree );
  free( backup_tree_fat );


  return 1;

}

void buildNeighborhoods(const char * atom_base, 
			struct store * neighStore, 
			struct mTree * neighTree, 

			size_t neighbor_count, 
			double radius, 
			
			void(*constructor)(const char* pdbCode, 
					   void * neigh,
					   struct store * atomStore, 
					   off_t center,
					   off_t neighbors[],
					   short count),
			int (*can_calculate_neighborhood)(struct store * s, off_t i)
			) {
  
  char * atom_fat, * atom_tree, * atom_tree_fat;
  struct store * atomStore;
  struct mTree * atomTree;
  struct atom * a;

  off_t i = 1;
  off_t start, stop;

  printf("%s ", atom_base);
  printf("(%.4s)\n", &(atom_base[strlen(atom_base)-10]));
  fflush(stdout);

  atom_fat = malloc(sizeof(char) * (strlen(atom_base) + 4 + 1));
  atom_tree = malloc(sizeof(char) * (strlen(atom_base) + 6 + 1));
  atom_tree_fat = malloc(sizeof(char) * (strlen(atom_base) + 10 + 1));

  strcpy(atom_fat,      atom_base);
  strcpy(atom_tree,     atom_base);
  strcpy(atom_tree_fat, atom_base);

  strcat(atom_fat,      ".fat");
  strcat(atom_tree,     ".mTree");
  strcat(atom_tree_fat, ".mTree.fat");

  atomStore  = store_open_cached(atom_base, atom_fat, sizeof(struct atom), 100, 100000);
  atomTree   = mTree_open_cached(atomStore, atom_tree, atom_tree_fat, 80000, (double(*)(const void*,const void*))&atom_distance);


  a = store_read(atomStore, i);
  

  //  atom_print(a);

  //  printf("%d %.4s :(\n", i, a->name);

  
  while(a->name[0] != 0) {

    int can_calc;
    
    //    int residue = a->residue;

    store_release(atomStore, i);

    //    printf("%.4s ", a->name);
    
    can_calc = can_calculate_neighborhood(atomStore, i);
    if(can_calc) {
      //      printf("%d\n", residue);
      neighborhood_indexer_calculate_neighborhood(&(atom_base[strlen(atom_base)-10]), radius, i, atomStore, atomTree, neighStore, neighTree, constructor, neighbor_count);
    }
    
    find_aa_range(atomStore, i, &start, &stop);
    
    i = stop + 1;
    
    a = store_read(atomStore, i);
    
  }

  store_release(atomStore, i);
  mTree_close(atomTree);
  store_close(atomStore);
  
  if(i % 10 == 0) {
    printf("Neighborhood tree stats:\n");
    mTree_printStats(neighTree);
    
  }

  free(atom_fat);
  free(atom_tree);
  free(atom_tree_fat);

}
/**


  @param pdbCode     The pdb code of the current structure.  (four characters)
  @param radius      The maximum distance between the center and one of its neighbors
  @param atomIndex   The position of the atom in atomStore
  @param atomStore   A store containing atoms.
  @param atomTree    An mTree that contains the atoms in atomStore, indexed by physical position.
  @param neighStore  The neighborhood store that new neighborhoods will be inserted into.
  @param neighTree   The tree that corresponds to neighStore
  
  @param constructor A function pointer.  It takes an array
  'neighbors[]' of length 'neigh_count' of offsets in the store
  'atomStore'.  It also requires the four character PDB code of the
  current PDB file.

  @param neigh_count The number of neighbors in this neighborhood.
  Does not include the center amino acid.  This is the length of the
  neighbors[] array that is passed into constructor.

  @todo This needs to be re-written so that the constructor returns
  true if it succeeded, false otherwise.  This function will then need
  to take a size parameter so that it can allocate some temporary
  space for the constructor to store its data in.   

*/
void neighborhood_indexer_calculate_neighborhood(const char * pdbCode, 
						 double radius,
						 off_t atomIndex,
						 struct store * atomStore, 
						 struct mTree * atomTree,
						 struct store * neighStore, 
						 struct mTree * neighTree,
						 void(*constructor)(const char* pdbCode, void * neigh,
								    struct store * atomStore, off_t center,
								    off_t neighbors[], 
								    short count),
						 short neigh_count
						 ) {


  off_t *  neighbor_indices  = malloc(sizeof(off_t) * neigh_count);
  
  neighborhood_indexer_find_neighbors(atomTree, atomStore, atomIndex, 
				      radius, neigh_count, neighbor_indices);
  
  if(neighbor_indices[neigh_count-1] != -1) {
    off_t n_idx = store_new(neighStore);
    void * neigh = store_read(neighStore, n_idx);
    
    constructor(pdbCode, neigh, atomStore, atomIndex, neighbor_indices, neigh_count);
    
    store_write(neighStore, n_idx);
  
    mTree_insert(neighTree, n_idx);
  
    store_release(neighStore, n_idx);
  }

  
  free (neighbor_indices);

}

void neighborhood_indexer_find_neighbors(struct mTree * tree, struct store * s, const off_t atomIndex, 
					 const double radius, const short neigh_count, off_t * neighbor_indices) {
  //  void aa_neighborhood_find_neighbors(struct mTree * tree, const struct store * s, struct atom * center,  double radius,
  //				      struct aa_neighborhood * neigh, short count) {

    // Initialize mTree query handler

  struct atom * center = store_read(s, atomIndex);
  
  struct knn_state * knn_state = knn_state_alloc(neigh_count, center, (double(*)(const void *, const void *))atom_distance, sizeof(struct atom));
  struct filter_state filter_state;
  short c;
  
  //  knn_state_reset(knn_state);
  
  filter_state.count        = 3;
  filter_state.functions    = malloc(sizeof(void *(*)(void *, void *)) * 3);
  filter_state.states       = malloc(sizeof(void *) * 3);
    //  filter_state.functions[0] = (void *(*)(void*, void*))filter_atom_by_name;                                                                                                                                         
  filter_state.functions[0] = (void *(*)(void*, void*))filter_atom_by_name;
  filter_state.states[0]    = " CA ";
  filter_state.functions[1] = (void *(*)(void*, void*))filter_atom_by_residue_number;
  filter_state.states[1]    = (void *)center->residue;
  filter_state.functions[2] = (void *(*)(void*, void*))knn;
  filter_state.states[2]    = knn_state;
  
  
  // Run the query.
  
  mTree_rangeQuery(tree, center, radius, (void(*)(void*,void*))filterAnd, &filter_state);
  
  for(c = 0; c < neigh_count; c++) {
    
    if(knn_state->results_distance[c] != MAXDOUBLE) {
      
      neighbor_indices[c] = ((struct atom*)knn_state->results)[c].serialNumber;
      
    } else {
      neighbor_indices[c] = -1;
    }
    
  }

  knn_state_free(knn_state);
  free(filter_state.functions);
  free(filter_state.states);
  store_release(s, atomIndex);
  
}



			//void daliesque_constructor(struct daliesque * d, struct xyz_int * center, unsigned char center_type_index, struct xyz_int neighbors[COUNT], unsigned char neighbor_type_indices[COUNT]));



int queryStructures(int argc, char * argv[],
		    void(*constructor)(const char* pdbCode, 
				       void * neigh,
				       struct store * atomStore, 
				       off_t center,
				       off_t neighbors[],
				       short count),
		    int (*can_calculate_neighborhood)(struct store * s, off_t i),
		    const void*(*print_fcn)(const void *, void *),
		    double (*distance)(const void * , const void *),
		             
		    size_t neighborhood_size,
		    short neighbor_count,
		    double radius,
		    double query_radius
		    ) {

  struct store * neigh_store;
  struct mTree * neigh_tree;
  char * atom_base;

  struct timeb ts;
  time_t start, stop;
  int last;
  off_t i, node_store_size;
  long * n;
  
  //  void openMainIndex(const char * main_neigh_base, struct store * neighStore, struct mTree * neighTree, size_t neigh_struct_size,
  //		     double(*neigh_distance)(const void*,const void*), size_t store_cache_size, size_t tree_cache_size);

  if(argc != 2) {
    printf("Usage: %s neighborhood_store\n\nPipe the query atom store base names into STDIN.\n\nsizeof(neighborhood) = %ld, sizeof(mTreeNode) = %ld\n\n", 
           argv[0], (long) neighborhood_size, (long) sizeof(struct mTreeNode));
    return 1;
  }

  openMainIndex(argv[1], &neigh_store, &neigh_tree, neighborhood_size, distance, 10000000, 8000000);

  //  ftime(&ts);

  //  last_sync = ts.time;

  ftime(&ts);
  start = ts.time;
  
  printf("Pre-reading tree nodes..\n");
  last = 100;

  i = 0;
  //  while(tmpNode->children[0] != 0) {                                                                                                                                                                               
  while(last) {
    struct mTreeNode * tmpNode;

    i++;
    tmpNode = store_read(neigh_tree->nStore, i);
    last --;
    if(tmpNode->children[0] != 0) {
      last = 100;
    }
    store_release(neigh_tree->nStore, i);
  }
  
  ftime(&ts);
  stop = ts.time;

  printf("..Done reading tree nodes.  # = %ld Time = %ld\n", (long int)i, (long int)stop - start);
  
  node_store_size = i;
  
  ftime(&ts);
  start = ts.time;
  
  printf("Pre-reading objects..\n");
  i =0;
  n = store_read(neigh_tree->oStore, i);
  last = 100;
  while(last && i < 10000000) {

    if(n[0] == 0 && n[1] == 0 && n[2] == 0) {
      last--;
    } else {
      last = 100;
    }
    
    store_release(neigh_tree->oStore, i);
    i++;
    n = store_read(neigh_tree->oStore, i);
  }
  store_release(neigh_tree->oStore, i);
  
  ftime(&ts);
  stop = ts.time;
  
  printf("..Done reading objects.  # = %ld (of size = %ld) time = %ld\n", (long int)i, (long int)neighborhood_size, (long int)stop-start);
  i = 0;
  
  

  while(EOF != scanf("%as\n", &atom_base)) {
    
    struct store * tmp_store;
    struct mTree * tmp_tree;

    if(unlink("tmp_store")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }
    if(unlink("tmp_store.fat")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }
    if(unlink("tmp_store.mTree")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }
    if(unlink("tmp_store.mTree.fat")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }

    openMainIndex("tmp_store", &tmp_store, &tmp_tree, neighborhood_size, distance, 12000, 10000);

    printf("Building index for %s...", atom_base);

    fflush(stdout);

    buildNeighborhoods(atom_base, tmp_store, tmp_tree, neighbor_count, radius, constructor, can_calculate_neighborhood);

    printf("done.\n");
    fflush(stdout);

    n = store_read(tmp_store, i);



    while(n[0] != 0 || n[1] != 0 || n[2] != 0) {
      off_t hits = 0;

      struct print_with_distance_state s;
      ftime(&ts);
      start = ts.time;

      

      printf("Q%ld:", (long)i);
      print_fcn(n, 0);
      printf("\n");

      s.query_object = n;
      s.distance = distance;
      s.print = print_fcn;
      s.print_state1 = 0;
      s.print_state2 = 0;

      //    aa_neighborhood_print(n);                                                                                                                                                                                    

      //    mTree_rangeQuery(db_mTree, n, atoi(argv[1]), (void(*)(void*,void*))&aa_neighborhood_print_hit, 0);                                                                                                           
      //mTree_rangeQuery(neigh_tree, n, query_radius, (void(*)(void*,void*))print_fcn, &hits);
      mTree_rangeQuery(neigh_tree, n, query_radius, (void(*)(void*,void*))print_with_distance, &s);
      store_release(tmp_store, i);
      ftime(&ts);
      stop = ts.time;

      printf("Query completed.  Time = %ld #hits = %ld\n", (long)stop-start, (long)hits);

      //i = -1;                                                                                                                                                                                                          
      i++;
      n = store_read(tmp_store,i);
    }

    store_release(tmp_store,i);

    mTree_close(tmp_tree);
    store_close(tmp_store);

    if(unlink("tmp_store")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }
    if(unlink("tmp_store.fat")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }
    if(unlink("tmp_store.mTree")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }
    if(unlink("tmp_store.mTree.fat")) { if (errno != ENOENT) { perror ("Couldn't unlink temporary file:"); return 1; } }

    
  }
  return 0;
}



int index_many_neighborhoods(int argc, char * argv[], 
			      void(*constructor)(const char* pdbCode, 
						 void * neigh,
						 struct store * atomStore, 
						 off_t center,
						 off_t neighbors[],
						 short count),
			      int (*can_calculate_neighborhood)(struct store * s, off_t i),
			      double (*distance)(const void * , const void *),
			      size_t neighborhood_size,
			      short neighbor_count,
			      double radius
			      ) {
  struct store * neigh_store;
  struct mTree * neigh_tree;
  char * atom_base;
  short first = 0;
  struct timeb ts;
  time_t last_sync, sync_interval;
  
  if(argc != 3 && argc != 4) {
    printf("Usage: %s neighborhood_store periodic_backup_base_name [recovery_filename]\n\nPipe the atom store base names into STDIN.\n\nThe recovery_filename option is the filename output by the last successful periodic backup.\nTo recover from a crash, copy the last good incremental backup to location of neighborhood\nstore, and pass in a recovery filename.\n\nsizeof(neighborhood) = %ld, sizeof(mTreeNode) = %ld\n\n", 
           argv[0], (long) neighborhood_size, (long) sizeof(struct mTreeNode));
    //  printf("Using these files for tree: %s, %s, %s, %s  sizeof(aa_neighborhood) = %ld\n", neigh_base, neigh_fat, neigh_tree, neigh_tree_fat, (long)sizeof(struct aa_neighborhood));
    return 1;
  }
  
  //  void openMainIndex(const char * main_neigh_base, struct store * neighStore, struct mTree * neighTree, size_t neigh_struct_size,
  //		     double(*neigh_distance)(const void*,const void*), size_t store_cache_size, size_t tree_cache_size);


  openMainIndex(argv[1], &neigh_store, &neigh_tree, neighborhood_size, distance, 10000000, 8000000);

  ftime(&ts);

  last_sync = ts.time;

  sync_interval = 7200;
  
  if(argc == 4) {
    while(EOF != scanf("%as\n", &atom_base)) {
      if (!strcmp(atom_base, argv[3])) {
        break;
      }
      printf("skipping %s\n", atom_base);
    }
    first = 1; // first = 1 => atom_base already contains the first file name.
  }
  

  while(first || EOF != scanf("%as\n", &atom_base)) {

    first = 0;

    if(!writeBackup(neigh_store, neigh_tree, atom_base, argv[2], &last_sync, sync_interval)) {
      assert(0);
    }
    
    buildNeighborhoods(atom_base, neigh_store, neigh_tree, neighbor_count, radius, constructor, can_calculate_neighborhood);
  }
  return 0;
  
  
}
